# flask_app.py
from flask import Flask, request, jsonify
import pandas as pd
import joblib
import torch
from pathformer.PathFormer import PathFormer, ModelConfig
from utils.predict_utils import predict_time_range

app = Flask(__name__)
config = ModelConfig()

# 全局版本变量
version = "v3.0.0"

# 设置参数


# 提取文件路径进行管理
data_path = 'dataset/ZGGG.09.02.2023.09.02.2025.1.0.0.cn.utf8.00000000.xls'
encoder_dd_path = f'model_files/label_encoder/label_encoder_dd_{version}.pkl'
encoder_c_path = f'model_files/label_encoder/label_encoder_c_{version}.pkl'
scaler_path = f'model_files/data_scaler/data_scaler_{version}.pkl'
model_dict_path = f'model_files/model_dict/pathformer_dict_{version}.pkl'

# 全局加载模型和编码器（启动时加载一次）
# encoder_dd = joblib.load(encoder_dd_path)
# encoder_c = joblib.load(encoder_c_path)
scaler = joblib.load(scaler_path)

# 在全局加载一次训练数据文件
data = pd.read_excel(data_path, skiprows=6)

# 初始化模型
model = PathFormer(
    input_dim=15,
    patch_size=48,
    embed_dim=256,
    num_heads=8,
    num_layers=6,
    output_dim=15,
    ff_dim=1024
)
# 加载模型参数字典
model.load_state_dict(torch.load(model_dict_path))
# 设备配置并加载到设备上
device = torch.device('cpu')
model.to(device)
model.eval()    # 设置模型为评估模式


def add_time_features(df):
    df['hour_sin'] = np.sin(2 * np.pi * df.index.hour / 24)
    df['hour_cos'] = np.cos(2 * np.pi * df.index.hour / 24)
    df['month_sin'] = np.sin(2 * np.pi * df.index.month / 12)
    df['month_cos'] = np.cos(2 * np.pi * df.index.month / 12)
    df['dayofweek_sin'] = np.sin(2 * np.pi * df.index.dayofweek / 7)
    df['dayofweek_cos'] = np.cos(2 * np.pi * df.index.dayofweek / 7)
    return df


def preprocess_data(original_data: pd.DataFrame) -> pd.DataFrame:
    """数据预处理（需与训练时一致）"""
    original_data.drop(columns=['ff10', 'WW', "W'W'"], inplace=True)
    original_data['当地时间 广州(机场)'] = pd.to_datetime(original_data['当地时间 广州(机场)'], format='%d.%m.%Y %H:%M')
    original_data.set_index('当地时间 广州(机场)', inplace=True)
    original_data.sort_index(ascending=True, inplace=True)

    original_data = add_time_features(original_data)

    original_data['VV'] = original_data['VV'].astype(str)
    original_data['VV'] = original_data['VV'].str.extract(r'(\d+\.\d+)').astype(float)
    original_data['DD'] = encoder_dd.transform(original_data[['DD']])
    original_data['c'] = encoder_c.transform(original_data[['c']])
    original_data.ffill(inplace=True)

    return original_data


def set_seed(seed):
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    random.seed(seed)


@app.route('/predict', methods=['POST'])
def predict():

    # 设置随机种子
    set_seed(42)

    # 解析请求参数
    req = request.json
    start_time = req.get('start_time')
    num_steps = min(int(req.get('num_steps', 10)), 336)  # 限制最大336步

    try:
        # 加载默认训练数据（示例需替换为实际数据源）
        # 首先将读取到的原始数据复制一份
        original_data = data.copy()
        processed_data = preprocess_data(original_data)

        # 执行预测
        pred_results = predict_time_range(
            model, scaler, processed_data, start_time, num_steps, seq_length=672
        )

        # 反编码结果
        # 将 Series 转换为 DataFrame
        pred_results['DD'] = encoder_dd.inverse_transform(pred_results[['DD']].astype(int)).flatten()
        pred_results['c'] = encoder_c.inverse_transform(pred_results[['c']].astype(int)).flatten()

        # 转换为可序列化的字典
        pred_results.index = pred_results.index.strftime('%Y-%m-%d %H:%M:%S')
        pred_results.rename(columns={pred_results.index.name: 'time'}, inplace=True)
        pred_results = pred_results.iloc[:, :9]

        # 将时间列单独提取出来，并作为每个记录中的一个键值对
        data_list = []
        for index, row in pred_results.iterrows():
            record = row.to_dict()
            record['time'] = index  # 添加时间键值对
            data_list.append(record)

        return jsonify({"status": "success", "data": data_list})

    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 400


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)
